# -*- coding: utf-8 -*-
from pandas import DataFrame
from numpy import nan

from core.dataClasses import StockTradeDataColumnName
from core.index_base import IndexBase

BREAK_POINT_HIGH_COLUMN = 'high_point'
BREAK_POINT_LOW_COLUMN = 'low_point'

def fun(x):
    return 1


class BreakPoint(IndexBase):
    def __init__(self):
        super(BreakPoint, self).__init__()

    def compute(self, data: DataFrame):
        data[BREAK_POINT_HIGH_COLUMN] = nan
        data[BREAK_POINT_LOW_COLUMN] = nan
        self.added_keys.extend([BREAK_POINT_HIGH_COLUMN, BREAK_POINT_LOW_COLUMN])

        data[BREAK_POINT_HIGH_COLUMN] = data[
            (data[StockTradeDataColumnName.HIGH] > data[StockTradeDataColumnName.HIGH].shift(1))
            & (data[StockTradeDataColumnName.HIGH] > data[StockTradeDataColumnName.HIGH].shift(2))
            & (data[StockTradeDataColumnName.HIGH] > data[StockTradeDataColumnName.HIGH].shift(3))
            & (data[StockTradeDataColumnName.HIGH] > data[StockTradeDataColumnName.HIGH].shift(4))
            & (data[StockTradeDataColumnName.HIGH] >= data[StockTradeDataColumnName.HIGH].shift(-1))
            & (data[StockTradeDataColumnName.HIGH] >= data[StockTradeDataColumnName.HIGH].shift(-2))
            & (data[StockTradeDataColumnName.HIGH] >= data[StockTradeDataColumnName.HIGH].shift(-3))
            & (data[StockTradeDataColumnName.HIGH] >= data[StockTradeDataColumnName.HIGH].shift(-4))
        ][BREAK_POINT_HIGH_COLUMN].replace(to_replace=nan, value=True)

        data[BREAK_POINT_LOW_COLUMN] = data[
            (data[StockTradeDataColumnName.LOW] < data[StockTradeDataColumnName.LOW].shift(1))
            & (data[StockTradeDataColumnName.LOW] < data[StockTradeDataColumnName.LOW].shift(2))
            & (data[StockTradeDataColumnName.LOW] < data[StockTradeDataColumnName.LOW].shift(3))
            & (data[StockTradeDataColumnName.LOW] < data[StockTradeDataColumnName.LOW].shift(4))
            & (data[StockTradeDataColumnName.LOW] <= data[StockTradeDataColumnName.LOW].shift(-1))
            & (data[StockTradeDataColumnName.LOW] <= data[StockTradeDataColumnName.LOW].shift(-2))
            & (data[StockTradeDataColumnName.LOW] <= data[StockTradeDataColumnName.LOW].shift(-3))
            & (data[StockTradeDataColumnName.LOW] <= data[StockTradeDataColumnName.LOW].shift(-4))
            # & (data[BREAK_POINT_HIGH_COLUMN] != True)
            # & (data[BREAK_POINT_HIGH_COLUMN].shift(1) != True)
            # & (data[BREAK_POINT_HIGH_COLUMN].shift(-1) != True)
            # & (data[BREAK_POINT_HIGH_COLUMN].shift(2) != True)
            # & (data[BREAK_POINT_HIGH_COLUMN].shift(-2) != True)
            # & (data[BREAK_POINT_HIGH_COLUMN].shift(3) != True)
            # & (data[BREAK_POINT_HIGH_COLUMN].shift(-3) != True)
            # & (data[BREAK_POINT_HIGH_COLUMN].shift(4) != True)
            # & (data[BREAK_POINT_HIGH_COLUMN].shift(-4) != True)
            ][BREAK_POINT_LOW_COLUMN].replace(to_replace=nan, value=True)
        data[[BREAK_POINT_LOW_COLUMN, BREAK_POINT_HIGH_COLUMN]] = data[[BREAK_POINT_LOW_COLUMN, BREAK_POINT_HIGH_COLUMN]].rolling(window=10, center=True).apply(fun, raw=True)
        print(data)


if __name__ == "__main__":
    from domain.transaction_data.repository import transaction_data_repository

    BreakPoint().compute(transaction_data_repository.get_contain_latest_data_by_tushare('002044.SZ'))